// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gtest/gtest.h>
#include <pollux/testing/expression/decimal_arg_types_generator_base.h>
#include <tests/expression/fuzzer/arg_types_generator_test_utils.h>

namespace kumo::pollux::fuzzer::test {

class DecimalArgTypesGeneratorTest : public testing::Test {
 protected:
  class UnaryArgTypesGenerator : public DecimalArgTypesGeneratorBase {
   public:
    UnaryArgTypesGenerator() {
      initialize(1);
    }

   protected:
    std::optional<std::pair<int, int>> toReturnType(int p, int s) override {
      auto precision = std::min(38, p + s + 1);
      auto scale = std::min(s + 1, 18);
      return {{precision, scale}};
    }
  };

  class BinaryArgTypesGenerator : public DecimalArgTypesGeneratorBase {
   public:
    BinaryArgTypesGenerator() {
      initialize(2);
    }

   protected:
    std::optional<std::pair<int, int>>
    toReturnType(int p1, int s1, int p2, int s2) override {
      auto s = std::max(s1, s2);
      auto p = std::min(38, std::max(p1 - s1, p2 - s2) + std::max(s1, s2) + 1);
      return {{p, s}};
    }
  };
};

TEST_F(DecimalArgTypesGeneratorTest, unary) {
  auto signature =
      exec::FunctionSignatureBuilder()
          .integerVariable("scale")
          .integerVariable("precision")
          .integerVariable("r_precision", "min(38, precision + scale + 1)")
          .integerVariable("r_scale", "min(scale + 1, 18)")
          .returnType("decimal(r_precision, r_scale)")
          .argumentType("decimal(precision, scale)")
          .build();

  const auto generator = std::make_shared<UnaryArgTypesGenerator>();
  for (auto returnType : {DECIMAL(10, 2), DECIMAL(38, 18)}) {
    assertReturnType(generator, *signature, returnType);
  }
  assertEmptyArgs(generator, *signature, DECIMAL(38, 20));
}

TEST_F(DecimalArgTypesGeneratorTest, binary) {
  auto signature =
      exec::FunctionSignatureBuilder()
          .integerVariable("a_scale")
          .integerVariable("b_scale")
          .integerVariable("a_precision")
          .integerVariable("b_precision")
          .integerVariable(
              "r_precision",
              "min(38, max(a_precision - a_scale, b_precision - b_scale) + max(a_scale, b_scale) + 1)")
          .integerVariable("r_scale", "max(a_scale, b_scale)")
          .returnType("decimal(r_precision, r_scale)")
          .argumentType("decimal(a_precision, a_scale)")
          .argumentType("decimal(b_precision, b_scale)")
          .build();

  const auto generator = std::make_shared<BinaryArgTypesGenerator>();
  for (auto returnType :
       {DECIMAL(10, 2), DECIMAL(38, 20), DECIMAL(38, 38), DECIMAL(38, 0)}) {
    assertReturnType(generator, *signature, returnType);
  }
}

} // namespace kumo::pollux::fuzzer::test
